package com.csthink.limiter.manager;

import com.google.common.collect.Maps;
import com.google.common.util.concurrent.RateLimiter;
import lombok.extern.slf4j.Slf4j;

import java.util.Map;
import java.util.concurrent.TimeUnit;

/**
 * 基于 Guava 的 RateLimiter 的限流实现 {@link RateLimiter}
 * @author <a href="mailto:security.2009@live.cn">Mars</a>
 * @since 2023-09-08
 */
@Slf4j
public class GuavaLimiter implements LimiterManager {

    private final Map<String, RateLimiter> limiterMap = Maps.newConcurrentMap();

    @Override
    public boolean tryAccess(Limiter limiter) {
        RateLimiter rateLimiter = getRateLimiter(limiter);
        if (rateLimiter == null) {
            return false;
        }

        // 尝试获取令牌
        boolean access = rateLimiter.tryAcquire(limiter.getTimeout(), TimeUnit.MILLISECONDS);
        log.info("Use guava limiter, {} access result :{}", limiter.getKey(), access);
        return access;
    }

    private RateLimiter getRateLimiter(Limiter limiter) {
        RateLimiter rateLimiter;
        // 验证缓存是否有命中key
        String key = limiter.getKey();
        if (!limiterMap.containsKey(key)) {
            // 创建令牌桶
            rateLimiter = RateLimiter.create(limiter.getPermitsPerSecond());
            limiterMap.put(key, rateLimiter);
            log.info("新建了令牌桶={}，容量={}", key, limiter.getPermitsPerSecond());
        }

        rateLimiter = limiterMap.get(key);
        return rateLimiter;
    }
}
